#include <cassert>
#include <vector>
#include <iostream>
#include <ctime>

#include "../core/i_rand.h"
#include "../geometry/i_plane.h"
#include "../geometry/i_util.h"
#include "i_ground.h"

namespace idl
{
    void ADBinocularGroundParam::set_default() {
        //ratio height*skyline = row index of upper image region, assume image region above skyline is mostly sky
        skyline = 0.55f;
        thre_gpf_inliers = 64;
        thre_gpf_sample_min = 128;
        thre_gpf_sample_max = 300;
        thre_gpf_dist = 150.f;     //unit in mm, 15cm = 0.15m
        thre_gpf_ransac_iter = 30; //max ransac 30 iterations
        thre_gpf_z_max = 20000.f;  //unit in mm, 20m
        thre_gpf_tilt_angle = 6.0f;
        gpf_roi_ignored_width = 4;
    }

    ADBinocularGroundTracker::ADBinocularGroundTracker(int width, int height, int track_length,
        const ADCameraIntrinsic<float>& intrinsic) : _width(width), _height(height) {
        assert(track_length > 0);
        _threeds = i_alloc_aligned<float>(_param.thre_gpf_sample_max * 4, 4);
        //2^4 = 16-byte aligned memory
        _indices = i_alloc2<int>(_param.thre_gpf_ransac_iter, 3);
        _hypothesis.resize(_param.thre_gpf_ransac_iter);
        _rand_seed = I_DEFAULT_SEED;
        _camera_intrinsic = intrinsic;
        _tracker = new ADGroundTrackManager(track_length);
        assert(_tracker != NULL);
    }

    ADBinocularGroundTracker::~ADBinocularGroundTracker() {
        i_free2<int>(_indices);
        i_free_aligned<float>(_threeds);
        if (_tracker) {
            delete _tracker;
        }
        _hypothesis.clear();
    }

    int ADBinocularGroundTracker::track(const unsigned char *const *disp, 
        const std::vector<float>& depthlut,
        float params[4]) {
        int skyline = i_round(_height * _param.skyline);
        int roi[] = { _param.gpf_roi_ignored_width, skyline,
        _width - (_param.gpf_roi_ignored_width << 1),
        _height - skyline };
        ADGroundPlane plane;
        int nr_inlier = detect(disp, depthlut, plane, roi, _width, _height);
        _tracker->push(plane);
        _tracker->get_ground(params);
        return (nr_inlier);
    }

    void ADBinocularGroundTracker::restart() {
        _tracker->restart();
    }

    int ADBinocularGroundTracker::detect(const unsigned char *const *disp, 
        const std::vector<float>& depthlut,
        ADGroundPlane &plane,
        const int roi[4], int width, int height) {
        int i = 0; 
        int j = 0; 
        int r = 0; 
        int c = 0; 
        int d = 0; 
        int pos = 0; 
        int step_x = 0; 
        int step_y = 0; 
        int nr_inlier = 0; 
        int nr_sample = 0;
        int nr_inlier_best = 0;
        int best = -1;
        int nr_disp = (int)depthlut.size();
        float *p_src = NULL; 
        float *p_dst = NULL;
        const unsigned char *disp_line = NULL;
        int roi_x = roi[0];
        int roi_y = roi[1];
        int roi_w = roi[2];
        int roi_h = roi[3];
        assert(roi_x >= 0 && roi_y >= 0 && roi_w <= width && roi_h <= height);
        assert(roi_w > 0 && roi_x + roi_w <= width);
        assert(roi_h > 0 && roi_y + roi_h <= height);
        //nr_sample_x * nr_sample_y <= thre_gpf_sample_max
        //nr_sample_y / nr_sample_x = roi_h / roi_w;
        //nr_sample_y = nr_sample_x * roi_h / roi_w;
        //nr_sample_x * nr_sample_x * roi_h / roi_w <= thre_gpf_sample_max;
        //(nr_sample_x)^2 <= thre_gpf_sample_max * roi_w / roi_h;
        int nr_sample_x = (int)(i_sqrt(_param.thre_gpf_sample_max * roi_w / roi_h));
        int nr_sample_y = (int)(nr_sample_x * roi_h) / roi_w;
        assert(nr_sample_x >= 2 && nr_sample_y >= 2);
        assert(nr_sample_x * nr_sample_y <= _param.thre_gpf_sample_max);
        step_x = roi_w / (nr_sample_x - 1);
        step_y = roi_h / (nr_sample_y - 1);
        float pt_to_pl_dist = 0.f; 
        float uv[2]; 
        float gpf_samples[9];
        //sample 3D points for ground plane fitting:
        for (r = roi_y; r < (roi_y + roi_h); r += step_y) {
            disp_line = disp[r];
            uv[1] = (float)r;
            for (c = roi_x; c < (roi_x + roi_w); c += step_x) {
                d = (int)disp_line[c];
                if (!d || d >= nr_disp) { //invalid disparity value
                    continue;
                }
                uv[0] = (float)c;
                pos = (nr_sample << 2);
                //backproject to the world coordinates 
                i_backproject_canonical(uv, _camera_intrinsic.K, depthlut.at(d), _threeds + pos);
                if (_threeds[pos + 2] > _param.thre_gpf_z_max) {
                    continue;
                }
                nr_sample++;
            }
        }

        std::cout << "# of ground fitting samples: " << nr_sample << std::endl;

        //failed to estimate ground plane - too few valid samples
        if (nr_sample < _param.thre_gpf_sample_min) {
            plane.force_invalid();
            std::cout << "Failed: not enough valid samples for RANSAC" << nr_sample << std::endl;
            return (0);
        }

        //generate random samples:
        for (i = 0; i < _param.thre_gpf_ransac_iter; i++) {
            i_random_sample(_indices[i], 3, nr_sample, _rand_seed);
        }

        /*vote for plane hypothesis*/
        for (i = 0; i < _param.thre_gpf_ransac_iter; ++i) {
            nr_inlier = 0;
            i_copy3(_threeds + (_indices[i][0] << 2), gpf_samples);
            i_copy3(_threeds + (_indices[i][1] << 2), gpf_samples + 3);
            i_copy3(_threeds + (_indices[i][2] << 2), gpf_samples + 6);
            i_plane_fit_destroyed(gpf_samples, _hypothesis.at(i).params);
            
            /*force to be unit norm*/
            _hypothesis.at(i).force_unit_norm();
            
            /*check if the plane hypothesis has valid geometry*/
            if (_hypothesis.at(i).get_degree_normal_to_y() > _param.thre_gpf_tilt_angle) {
                continue;
            }

            /*iterate samples and check if the point to plane distance is within a threshold*/
            p_src = _threeds;
            for (j = 0; j < nr_sample; ++j, p_src += 4) {
                pt_to_pl_dist = i_plane_to_point_distance_w_unit_norm(
                    _hypothesis.at(i).params, p_src);
                if (pt_to_pl_dist < _param.thre_gpf_dist) {
                    nr_inlier++;
                }
            }

            _hypothesis.at(i).set_nr_support(nr_inlier); //record the # of supports
            
            //record the best plane index
            if (nr_inlier > nr_inlier_best) {
                nr_inlier_best = nr_inlier;
                best = i;
            }
        }

        /*if not pass the minimal number of inliers request or
        no valid geometry*/
        if (nr_inlier_best < _param.thre_gpf_inliers || best == -1) {
            plane.force_invalid();
            std::cout << "Failed: not enough inlier samples for RANSAC" << nr_sample << std::endl;
            return (0);
        }

        /*assign the best plane*/
        plane = _hypothesis.at(best);

        /*iterate samples and check if the point to plane distance is within threshold*/
        nr_inlier = 0;
        p_src = _threeds;
        p_dst = _threeds;
        for (j = 0; j < nr_sample; ++j, p_src += 4) {
            pt_to_pl_dist = i_plane_to_point_distance_w_unit_norm(plane.params, p_src);
            if (pt_to_pl_dist < _param.thre_gpf_dist) {
                i_copy3(p_src, p_dst);
                p_dst += 3;
                nr_inlier++;
            }
        }
        /*note that _pf_threeds will be destroyed after calling this routine*/
        i_plane_fit_total_least_square(_threeds, plane.params, nr_inlier);
        //assigne # of supports;
        plane.set_nr_support(nr_inlier);
        //force unit norm
        plane.force_unit_norm();
        //force positive y direction - upper direction
        plane.force_positive_normal_y();
        /*filtering: the best plane is not valid - geometry error*/
        if (plane.get_degree_normal_to_y() > _param.thre_gpf_tilt_angle) {
            plane.force_invalid();
            std::cout << "Failed: geometry error, not a valid ground plane" << std::endl;
            return (0);
        }

        return (nr_inlier);
    }

    ADGroundTrackManager::ADGroundTrackManager(int track_length) {
        assert(track_length > 0);
        _tracks.resize(track_length);
        _const_weight_temporal.resize(track_length, 0.f);
        _weight.resize(track_length, 0.f);
        for (int i = track_length - 1; i >= 0; i--) {
            _const_weight_temporal.at(i) = i_pow(i_sqrt(2.0f), track_length - 1 - i);
        }
        //normalize;
        float accm_wei = i_sum(_const_weight_temporal.data(), track_length);
        i_scale(_const_weight_temporal.data(), track_length, i_rec(accm_wei));
        _head = track_length; //init value 
    }

    ADGroundTrackManager::~ADGroundTrackManager() {
    }

    void ADGroundTrackManager::push(const ADGroundPlane &plane) {
        int i = 0; 
        int length = (int)_tracks.size();
        ADGroundPlaneSpherical ps;
        plane_eucli_to_spher(plane, ps);
        if (_head == 0) {
            if (length > 1) { //move backwards
                for (i = length - 1; i >= 1; i--) {
                    _tracks[i] = _tracks[i - 1];
                }
            }
        } else { //_head >=  zero
            _head = _head - 1;
        }
        _tracks[_head] = ps; //_head representing the current head position;
    }

    void ADGroundTrackManager::get_ground(float params[4]) {
        int i = 0;
        int length = (int)_tracks.size();
        float accm_wei = 0.f; 
        float w = 0.f; 
        float nr_support = 0.f;
        ADGroundPlaneSpherical ps;
        if (!length) {
            i_zero4(params);
            return;
        } else {
            if (length == 1) {
                plane_spher_to_eucli(_tracks.front(), params); // guaranteed a unit norm
                return;
            }
        }

        //i_scale(_weight.data() + _head, (length - _head), i_rec(accm_wei));
        for (i = _head; i < length; i++) {
            w = (float)_tracks.at(i).get_nr_support();//if not valid, nr_support = 0
            _weight.at(i) = _const_weight_temporal.at(i) * w;
        }
        
        //normalize factor;
        accm_wei = i_sum(_weight.data() + _head, (length - _head));
        
        for (i = _head; i < length; i++) {
            w = _weight.at(i);
            ps.theta += (_tracks.at(i).theta * w);
            ps.phi += (_tracks.at(i).phi * w);
            ps.d += (_tracks.at(i).d * w);
            nr_support += (_tracks.at(i).get_nr_support() * w);
        }

        ps.theta = i_div(ps.theta, accm_wei);
        ps.phi = i_div(ps.phi, accm_wei);
        ps.d = i_div(ps.d, accm_wei);
        nr_support = i_div(nr_support, accm_wei);
        ps.set_nr_support(i_round(nr_support));
        plane_spher_to_eucli(ps, params);
        return;
    }
    
    void ADGroundTrackManager::restart() {
        unsigned int i = 0; 
        for (; i < _tracks.size(); i++) {
            _tracks.at(i).force_invalid();
            _weight.at(i) = 0.f;
        }
        _head = (int)_tracks.size(); //reset to init value
    }

    //this routine assumes the surface normal has been normalized outside, therefore
    //{src.params[0], src.params[1], src.params[2]} is a unit vector
    void ADGroundTrackManager::plane_eucli_to_spher(const ADGroundPlane & src, 
        ADGroundPlaneSpherical & dst) {
        if (src.is_valid() == false) {
            dst.force_invalid();
        } else {
            ADGroundPlane p = src;
            assert(p.params[1] != 0 || p.params[0] != 0);
            p.force_unit_norm(); //to be safe
            p.force_positive_normal_y();
            dst.theta = i_acos(p.params[2]);
            dst.phi = i_atan2(p.params[0], p.params[1]);
            dst.d = p.params[3];
            dst.set_nr_support(src.get_nr_support());
        }
    }

    void ADGroundTrackManager::plane_spher_to_eucli(const ADGroundPlaneSpherical & src, 
        ADGroundPlane & dst) {
        if (src.is_valid() == false) {
            dst.force_invalid();
        } else {
            float nz = i_cos(src.theta);
            float nx_over_ny = i_tan(src.phi);//assume positive ny;
            //nx = ny * nx_over_ny;
            //nx * nx + ny * ny + nz * nz = 1.0;
            //ny^2 + ny^2*nx_over_ny^2 = 1 - nz^2;
            //ny^2(1 + nx_over_ny^2) = 1 - nz^2;
            //ny = sqrt((1-nz^2)/(1 + nx_over_ny^2))
            //ny is positive, guaranteed
            float ny = i_sqrt((1 - nz * nz) / (1 + nx_over_ny * nx_over_ny));
            float nx = ny * nx_over_ny;
            dst.params[0] = nx;
            dst.params[1] = ny;
            dst.params[2] = nz;
            dst.params[3] = src.d;
            dst.set_nr_support(src.get_nr_support());
        }
    }

    void ADGroundTrackManager::plane_spher_to_eucli(const ADGroundPlaneSpherical & src, 
        float params[4])
    {
        if (src.is_valid() == false) {
            i_zero4(params);
        } else {
            float nz = i_cos(src.theta);
            float nx_over_ny = i_tan(src.phi);//assume positive ny;
                                              //nx = ny * nx_over_ny;
                                              //nx * nx + ny * ny + nz * nz = 1.0;
                                              //ny^2 + ny^2*nx_over_ny^2 = 1 - nz^2;
                                              //ny^2(1 + nx_over_ny^2) = 1 - nz^2;
                                              //ny = sqrt((1-nz^2)/(1 + nx_over_ny^2))
                                              //ny is positive, guaranteed
            float ny = i_sqrt((1 - nz * nz) / (1 + nx_over_ny * nx_over_ny));
            float nx = ny * nx_over_ny;
            params[0] = nx;
            params[1] = ny;
            params[2] = nz;
            params[3] = src.d;
        }
    }
} //namespace idl